/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package hmm;

import java.math.*;
/**
 *
 * @author noe
 */
public class AlphaBeta {
    
    // The parameters model 
    private double[][] A, B;
    private double[] pi;
    
    // The matrices used to readjust the parameters
    private int[][] observations;
    private double [][][] Alpha, Beta;
    private double [][] Norm;
    
    
    // Number of states
    private static final int NUM_ST = 2;
    //The cardinality of the set of observations
    private static final int NUM_OBS = 3;
    
    // The number of observations sequences
    private int nObSeq;
    
    private double resu;
    
    // Field needed to compute the optimal solution
    // usign dynamic programing, i.e. Viterbi
    private double[][][] dp;    
    private int[][][] dpPath; 
    private int[][] optimalSeq;
    
    String seqDP;

    public AlphaBeta() 
    {
        /*this.A = new double[][]{{0.7,0.3},{0.4,0.6}}; // N = 2
        this.B = new double[][]{{0.5,0.3,0.2},{0.1,0.6,0.3}}; //M = 3
        this.pi = new double[]{0.6,0.4};*/
        this.pi = new double[]{0.5,0.5};
        this.A = new double[][]{{0.5,0.5},{0.5,0.5}}; // N = 2
        this.B = new double[][]{{0.4,0.1,0.5},{0.1,0.5,0.4}}; //M = 3
                
        this.observations = new int[][]{{2,0,0,2,1,2,1,1,1,2,1,1,1,1,1,2,2,0,0,1}};
        nObSeq = observations.length;
        this.Alpha = new double [nObSeq][][];
        this.Beta = new double [nObSeq][][];
        this.Norm = new double [nObSeq][];
        
        // Dynamic programming solutioin
        dp = new double[nObSeq][][];
        dpPath = new int[nObSeq][][];
        optimalSeq = new int[nObSeq][];
        
        resu = 0;
    }
    
    public void AlphaMatrix()
    {
        for (int k=0;k<nObSeq;k++)
        {
            int length = observations[k].length;
            Alpha[k] = new double [NUM_ST][length];
            Norm[k] = new double [length];
        }
        
        // We start with the base case for the forward variable
        double norm;
        
        // Initializing step
        for (int k=0;k<nObSeq;k++)
        {
            norm = 0;
            for (int i=0;i<NUM_ST;i++)
            {
                Alpha[k][i][0] = pi[i]*B[i][observations[k][0]];
                norm += Alpha[0][i][k];
            }
            
            // Scalling the values of A[0][i]
            Norm[k][0] = 1/norm;        
            for (int i=0;i<NUM_ST;i++)
                Alpha[k][i][0] *= Norm[k][0];
        }
        
        //
        // Computation of the forward variable for values of time greater than 0
        //
        for (int k=0;k<nObSeq;k++)
        {
            for (int t=1;t<observations[k].length;t++)
            {    
                norm = 0;
                for (int j=0;j<NUM_ST;j++)
                {
                    Alpha[k][j][t] = 0; // Initialization. Set values to zero

                    for (int i=0;i<NUM_ST;i++)
                        Alpha[k][j][t] += Alpha[k][i][t-1]*A[i][j];

                    Alpha[k][j][t] *= B[j][observations[k][t]];
                    norm += Alpha[k][j][t];
                }

                // Scalling process for Alpha[t][j]
                Norm[k][t] = 1/norm;
                for (int j=0;j<NUM_ST;j++)
                    Alpha[k][j][t] *= Norm[k][t];
            }
        }
    }
    
    public void BetaMatrix()
    {
        for (int k=0;k<nObSeq;k++)
            Beta[k] = new double [NUM_ST][observations[k].length];
        
        // Base case
        for (int k=0;k<nObSeq;k++)
        {
            int length = observations[k].length - 1;
            for (int i=0;i<NUM_ST;i++)
                Beta[k][i][length] = Norm[k][length];
        }
        
        //
        // Inductive state for the backward algorithm
        //
        for (int k=0;k<nObSeq;k++)
        {
            for (int t=observations[k].length-2; t>=0; t--)
                for (int i=0; i<NUM_ST;i++)
                {
                    Beta[k][i][t] = 0; //Initialization. Set values to zero

                    for (int j=0;j<NUM_ST;j++)
                        Beta[k][i][t] += A[i][j]*B[j][observations[k][t+1]]*Beta[k][j][t+1];

                    // It's time to scale the current value
                    Beta[k][i][t] *= Norm[k][t];
                }
        }
    }
    
    public double alpha(int t, int i, int k) //t= time, i=state, k = obs sequence
    {
        if(t == 0)
            return pi[i]*B[i][observations[k][0]];
        else
        {
            double sum = 0;
            for(int j=0; j<A.length; j++)
                sum += alpha(t-1, j,k)*A[j][i];
            
            return sum*B[i][observations[k][t]];
        }
    }
    
    public double alphaPass()
    {
        double sum = 0;
        int k = 0;
    
        for(int j=0; j<A.length; j++)
            sum += alpha(observations.length-1, j, k);
        
        return sum;
    }
    
    public double beta(int t, int i, int k) //t= time, i=state, k = obs sequence
    {
        if(t == observations[k].length-1)
            return 1;
        else
        {
            double sum = 0;
            for(int j = 0; j <NUM_ST; j++)
                sum += A[i][j]*B[j][observations[k][t+1]]*beta(t+1, j, k);
            
            return sum;
        }
    }
    
    public double betaPass()
    {
        /*double sum = 0, partialSum;
        
        for(int j=0; j<A.length;j++)
        {
            partialSum = pi[j]*B[j][observations[0]]*beta(0, j);
            
            System.out.println(String.format("Value from state " + String.valueOf(j) + " is: %1$.6f", partialSum));
            
            sum += partialSum;
        }*/
        double sum = 0;
        int k = 0;
        
        for(int j=0; j<A.length; j++)
            sum += pi[j]*B[j][observations[k][0]]*beta(0, j, k);//Notice the multiplication 
                                                          //with the initial distribution
        
        return sum;
    }
    
    public double gamma(int time, int i, int k)
    {       
        return alpha(time, i, k)*beta(time, i, k)/alphaPass();
    }
    
    public double xi(int time, int i, int j, int k)
    {
        return alpha(time,i,k)*A[i][j]*B[j][observations[k][time+1]]*beta(time+1,j,k)/alphaPass();
    }
    
    public void baulWelch()
    {
        int numIte = 20;
        
        while (numIte > 0)
        {
            // We update the matrices that store the values of alpha and beta
            AlphaMatrix();
            BetaMatrix();
        
            //----------------------            
            // Reestimating Pi
            //----------------------
            double numPi,denomPi,tempPi;
            
            for (int i=0; i<NUM_ST; i++)
            {                
                pi[i] = 0;
                
                for (int k=0;k<nObSeq;k++)
                {
                    numPi=0;
                    denomPi=0;
                    tempPi=0;
                
                    for (int n=0;n<NUM_ST;n++)
                    {
                        for (int m=0;m<NUM_ST;m++)
                        {
                            tempPi = Alpha[k][n][0]*A[n][m]*B[m][observations[k][1]]*Beta[k][m][1];
                            
                            if (i==n)
                               numPi += tempPi;
                            
                            denomPi += tempPi;                          
                        }                                     
                    }     
                    
                    pi[i] += numPi/denomPi;
                }
                
                pi[i] /= nObSeq;
            }
            
            
            //----------------------
            // Reestimating A
            //----------------------
            double gamma;
            double xi;
            
            for (int i=0;i<NUM_ST;i++)
            {
                gamma = 0;
                for (int j=0;j<NUM_ST;j++)
                {
                    xi = 0;
                    for (int k=0;k<nObSeq;k++)
                        for (int t=0;t<observations[k].length-1;t++)
                            xi += Alpha[k][i][t]*A[i][j]*B[j][observations[k][t+1]]*Beta[k][j][t+1];
                    
                    A[i][j] = xi;
                    gamma += xi;                    
                }
                
                for (int j=0;j<NUM_ST;j++)
                    A[i][j] /= gamma; 
            }
            
            
            //----------------------
            // Restimating B
            //----------------------
            /*for (int i=0;i<NUM_ST;i++)
                for (int j=0; j<B.length;j++)
                {
                    double gammaNum=0, gammaDen=0;
                    for (int t=0; t<observations.length;t++)
                    {
                        if(observations[t]==j)
                            gammaNum += gamma(t,i);
                        
                        gammaDen += gamma(t,i);
                    }
                    B[i][j] = gammaNum / gammaDen;
                }
            */
            double numBeta, denBeta;
            for (int i=0;i<NUM_ST;i++)
                for (int m=0; m<NUM_OBS;m++)
                {
                    numBeta=0;
                    denBeta=0;
                    
                    for (int k=0; k<nObSeq;k++)
                        for (int t=0;t<observations[k].length;t++)
                        {
                            if(observations[k][t]==m)
                                numBeta += Alpha[k][i][t]*Beta[k][i][t]; 

                            denBeta += Alpha[k][i][t]*Beta[k][i][t];
                        }
                    
                    B[i][m] = numBeta / denBeta;
                }
            
            numIte--;
        }
        
        String arr="", myPi="", arrB="";
        for (int i=0; i<NUM_ST;i++)
        {
            myPi += String.valueOf(pi[i]) + " ";
            
            for (int j=0;j<NUM_ST;j++)
                arr += String.valueOf(A[i][j]) + " ";
            
            for (int k=0; k<NUM_OBS;k++)
                arrB += String.valueOf(B[i][k]) + " ";
            
            arr += '\n';
            arrB += '\n';
        }
        
        /*double probVal = 1;
        for (int i=0;i<TIME;i++)
            probVal *= Norm[i];*/
       
        System.out.println("Pi\n" + myPi + "\nA\n" + arr + "\nB\n" + arrB);
    }
    
    public int getT()
    {
        return observations.length;
    }
    
    public int getN()
    {
        return A.length;
    }
    
    public void dynamicProgramming()
    {
        for (int k=0;k<nObSeq;k++)
        {
            int length = observations[k].length;
            dp[k] = new double [NUM_ST][length];
            dpPath[k] = new int [NUM_ST][length-1];
            optimalSeq[k] = new int[observations[k].length];
        }
        
        seqDP = "";
        
        int rep =2;
        while (rep > 0)
        {
            // We find the optimal sequence using viterbi and do statistical 
            // reasoning for all training sequences
            for (int k=0;k<observations.length;k++)
            {
                // We fill the first column of matrix dp for this particular k
                for(int i=0; i<NUM_ST;i++)
                    dp[k][i][0] = Math.log(pi[i])+Math.log(B[i][observations[k][0]]);
                
                //We fill the rest of matrices dp and dpPath
                //The way we do it is column by column
                for(int t=1; t<observations[k].length; t++)
                {
                    //This loop aims to fill the cell for each state in the column t 
                    for(int ns=0; ns<NUM_ST;ns++)
                    {
                        //The current state is ns
                        //To know the value at each position we need to compute the maximum value
                        double maxVal = Double.NEGATIVE_INFINITY, tmp;
                        int maxState = 0;
                        for(int i=0; i<NUM_ST;i++)
                        {
                            tmp = dp[k][i][t-1]+Math.log(A[i][ns]);
                            if(tmp > maxVal)
                            {
                                maxVal = tmp;
                                maxState = i;
                            }
                        }
                        dp[k][ns][t] = maxVal+Math.log(B[ns][observations[k][t]]);
                        dpPath[k][ns][t-1] = maxState;
                    }
                }
                
                //Now we want to find the final state with the highest probability
                //in the state sequence, so we analize the last column
                double max = Double.NEGATIVE_INFINITY;
                int state = 0;
                for(int i=0; i<NUM_ST;i++)
                {
                    if(dp[k][i][observations[k].length-1]>max)
                    {
                        max = dp[k][i][observations[k].length-1];
                        state = i;
                    }
                }
                
                //We do the following to know the path that leads us
                //to the optimal solution                
                String str = "]";
                optimalSeq[k][observations[k].length-1] = state;
                str = String.valueOf(state) + str;
                for(int t = observations[k].length-2;t>=0;t--)
                {
                    state = dpPath[k][state][t];
                    str = String.valueOf(state) + " " + str;
                    optimalSeq[k][t] = state;
                }
                
                seqDP = "["+ str;
            }
            
            A = new double[NUM_ST][NUM_ST];
            B = new double[NUM_ST][NUM_OBS];
            pi = new double[NUM_ST];
            for (int k=0;k<nObSeq;k++)
            {     
                int t;
                pi[optimalSeq[k][0]]++;
                for(t=0;t<observations[k].length-1;t++)
                {
                    A[optimalSeq[k][t]][optimalSeq[k][t+1]]++;
                    B[optimalSeq[k][t]][observations[k][t]]++;
                }
                
                B[optimalSeq[k][t]][observations[k][t]]++;
            }
            
            for(int i=0; i<NUM_ST;i++)
            {
                /////////////////////////////////
                // Recomputing matrix A
                /////////////////////////////////
                int totalA = 0;
                for (int j=0;j<NUM_ST;j++)
                    totalA += A[i][j];
                
                for (int j=0;j<NUM_ST;j++)
                {
                    System.out.print(A[i][j] + " ");
                    A[i][j] /= totalA;
                }
                
                System.out.println("");
                
                /////////////////////////////////
                // Recomputing matrix B
                /////////////////////////////////
                int totalB = 0;
                for (int m=0;m<NUM_OBS;m++)
                    totalB += B[i][m];
                
                for (int m=0;m<NUM_OBS;m++)
                {
                    System.out.print(B[i][m] + " ");
                    B[i][m] /= totalB;
                }
                
                System.out.println("");
                
                /////////////////////////////////
                // Recomputing initial distribution Pi
                /////////////////////////////////
                pi[i] /= nObSeq;
            }
            
            System.out.println("\n");
          
            rep--;
            
            System.out.println(seqDP);
        }
        
        ////////////////////////////
        // Done reestimating
        ////////////////////////////
        
        String arr="", myPi="", arrB="";
        for (int i=0; i<NUM_ST;i++)
        {
            myPi += String.valueOf(pi[i]) + " ";
            
            for (int j=0;j<NUM_ST;j++)
                arr += String.valueOf(A[i][j]) + " ";
            
            for (int k=0; k<NUM_OBS;k++)
                arrB += String.valueOf(B[i][k]) + " ";
            
            arr += '\n';
            arrB += '\n';
        }
        
        /*double probVal = 1;
         * for (int i=0;i<TIME;i++)
         * probVal *= Norm[i];*/
        
        System.out.println("Pi\n" + myPi + "\nA\n" + arr + "\nB\n" + arrB);
    }
    
    public String sequenceDP()
    {
        return seqDP;
    }
        
    /*public void combinations()
    {
        List<Number> list = new ArrayList<>();
        list.add(0);
        list.add(1);
        printCombinations(new ArrayList<Number>(), list, 0);
    }
    
    public void printCombinations(List<Number> done, List<Number> numbers, int depth)
    {
        double normalize = alphaPass();
        if (observations.length <= depth) {
            int ind = (int)done.get(0), indNext;
            double res = pi[ind]*B[ind][observations[0]];
            for(int i=0; i<observations.length-1;i++)
            {
                indNext = (int)done.get(i+1);
                res *= A[ind][indNext]*B[indNext][observations[i+1]];
                ind = indNext;
            }
            resu += res;
            System.out.println(String.format("P"+ done + " = %.6f", res/normalize));
            System.out.println(String.format("Current value = %.6f", resu/normalize));
        } else {
            for (Number r : numbers) {
                List<Number> newDone = new ArrayList<>(done);
                newDone.add(r);
                printCombinations(newDone, numbers, depth + 1);
            }
        }
    }*/
}
